# -*- coding:utf-8 -*-

from sqlmodel.ext.asyncio.session import AsyncSession
from aioredis import Redis
from contextlib import asynccontextmanager

from ...models import StorageMode, OrmModel
from ...engine import HybridEngine

class AsyncHybridEngine(HybridEngine):
    def __init__(self, sql_url: str, redis_url: str):
        super().__init__(sql_url=sql_url, redis_url=redis_url)
        # 覆盖为异步客户端
        self.redis = Redis.from_url(redis_url)

    @asynccontextmanager
    async def session(self):
        """异步上下文管理器"""
        async_session = AsyncHybridSession(self)
        try:
            yield async_session
            await async_session.commit()
        except Exception as e:
            await async_session.rollback()
            raise e
        finally:
            await async_session.close()

class AsyncHybridSession(AsyncSession):
    def __init__(self, engine: AsyncHybridEngine):
        # 调用父类的构造方法，传入SQL引擎
        super().__init__(engine.sql_engine)
        # 将传入的异步混合引擎赋值给实例变量self.engine
        self.engine = engine

    async def asave(self, instance: OrmModel):
        """异步保存"""
        # SQL 保存
        if instance.__storage_mode__ in [StorageMode.SQL, StorageMode.HYBRID]:
            self.add(instance)
            await self.flush()  # 获取自增ID

        # Redis 保存
        if instance.__storage_mode__ in [StorageMode.REDIS, StorageMode.HYBRID]:
            await self.engine.redis.hset(
                instance.redis_key(),
                mapping=instance.model_dump()
            )
            if instance.__redis_ttl__:
                await self.engine.redis.expire(instance.redis_key(), instance.__redis_ttl__)

    async def aget(self, model_cls: type[OrmModel], pk: str):
        """异步获取"""
    # 检查模型类的存储模式是否为REDIS
        if model_cls.__storage_mode__ == StorageMode.REDIS:
        # 从Redis中获取所有与模型类相关的哈希数据
            data = await self.engine.redis.hgetall(model_cls().redis_key())
        # 将获取到的二进制数据解码为字符串，并构造模型实例返回
            return model_cls(**{k.decode(): v.decode() for k, v in data.items()})
        else:
        # 如果存储模式不是REDIS，则调用父类的get方法获取数据
            return await super().get(model_cls, pk)